import os
import cv2
import math
import numbers
import random

from PIL import Image, ImageOps
import numpy as np

import torch
import torchvision.transforms as transforms
from torchvision.transforms import functional as F

# from infrastructure.photo_dl.infrastructure import utils


class Compose(object):
    def __init__(self, transforms):
        self.transforms = transforms

    def __call__(self, img, masks):
        for t in self.transforms:
            img, masks = t(img, masks)
        return img, masks

    def extend(self, transforms):
        self.transforms += transforms
        return self

    def __repr__(self):
        return str([type(t).__name__ for t in self.transforms])


class RandomCrop(object):
    def __init__(self, size, padding=0):
        if isinstance(size, numbers.Number):
            self.size = (int(size), int(size))
        else:
            self.size = size
        self.padding = padding

    def __call__(self, imgs, masks):
        if self.padding > 0:
            imgs = tuple(ImageOps.expand(img, border=self.padding, fill=0) for img in imgs)
            masks = tuple(ImageOps.expand(mask, border=self.padding, fill=0) for mask in masks)

        w, h = imgs[0].size
        th, tw = self.size
        if w == tw and h == th:
            return imgs, masks
        if w < tw or h < th:
            return tuple(img.resize((tw, th), Image.BILINEAR) for img in imgs), \
                   tuple(mask.resize((tw, th), Image.NEAREST) for mask in masks)

        x1 = random.randint(0, w - tw)
        y1 = random.randint(0, h - th)
        return tuple(img.crop((x1, y1, x1 + tw, y1 + th)) for img in imgs), \
               tuple(mask.crop((x1, y1, x1 + tw, y1 + th)) for mask in masks)


class CenterCrop(object):
    def __init__(self, size):
        if isinstance(size, numbers.Number):
            self.size = (int(size), int(size))
        else:
            self.size = size

    def __call__(self, img, mask):
        assert img.size == mask.size
        w, h = img.size
        th, tw = self.size
        x1 = int(round((w - tw) / 2.))
        y1 = int(round((h - th) / 2.))
        return img.crop((x1, y1, x1 + tw, y1 + th)), mask.crop((x1, y1, x1 + tw, y1 + th))


class RandomHorizontallyFlip(object):
    def __call__(self, imgs, masks):
        if random.random() < 0.5:
            return tuple(img.transpose(Image.FLIP_LEFT_RIGHT) for img in imgs), \
                   tuple(mask.transpose(Image.FLIP_LEFT_RIGHT) for mask in masks)
        return imgs, masks


class FreeScale(object):
    def __init__(self, size):
        self.size = tuple(reversed(size))  # size: (h, w)

    def __call__(self, imgs, masks):
        return tuple(img.resize(self.size, Image.BILINEAR) for img in imgs), \
               tuple(mask.resize(self.size, Image.NEAREST) for mask in masks)


class Scale(object):
    def __init__(self, size):
        self.size = size

    def __call__(self, img, mask):
        assert img.size == mask.size
        w, h = img.size
        if (w >= h and w == self.size) or (h >= w and h == self.size):
            return img, mask
        if w > h:
            ow = self.size
            oh = int(self.size * h / w)
            return img.resize((ow, oh), Image.BILINEAR), mask.resize((ow, oh), Image.NEAREST)
        else:
            oh = self.size
            ow = int(self.size * w / h)
            return img.resize((ow, oh), Image.BILINEAR), mask.resize((ow, oh), Image.NEAREST)


class RandomSizedCrop(object):
    def __init__(self, size, scale=(0.9, 1.1), padding=0):
        self.size = size
        self.padding = padding
        self.scale = scale

    def __call__(self, imgs, masks):
        if self.padding > 0:
            imgs = tuple(ImageOps.expand(img, border=self.padding, fill=0) for img in imgs)
            masks = tuple(ImageOps.expand(mask, border=self.padding, fill=0) for mask in masks)

        w, h = imgs[0].size
        th, tw = self.size
        scale = random.uniform(*self.scale)
        th, tw = round(th * scale), round(tw * scale)

        if w == tw and h == th:
            return imgs, masks
        if w < tw or h < th:
            return tuple(img.resize((tw, th), Image.BILINEAR) for img in imgs), \
                   tuple(mask.resize((tw, th), Image.NEAREST) for mask in masks)

        x1 = random.randint(0, w - tw)
        y1 = random.randint(0, h - th)
        return tuple(img.crop((x1, y1, x1 + tw, y1 + th)) for img in imgs), \
               tuple(mask.crop((x1, y1, x1 + tw, y1 + th)) for mask in masks)


class RandomRotate(object):
    def __init__(self, degree):
        self.degree = degree

    def __call__(self, imgs, masks):
        rotate_degree = random.random() * 2 * self.degree - self.degree

        return tuple(img.rotate(rotate_degree, Image.BILINEAR) for img in imgs), \
               tuple(mask.rotate(rotate_degree, Image.NEAREST, fillcolor=0) for mask in masks)


class RandomSized(object):
    def __init__(self, size):
        self.size = size
        self.scale = Scale(self.size)
        self.crop = RandomCrop(self.size)

    def __call__(self, img, mask):
        assert img.size == mask.size

        w = int(random.uniform(0.5, 2) * img.size[0])
        h = int(random.uniform(0.5, 2) * img.size[1])

        img, mask = img.resize((w, h), Image.BILINEAR), mask.resize((w, h), Image.NEAREST)

        return self.crop(*self.scale(img, mask))


class SlidingCropOld(object):
    def __init__(self, crop_size, stride_rate, ignore_label):
        self.crop_size = crop_size
        self.stride_rate = stride_rate
        self.ignore_label = ignore_label

    def _pad(self, img, mask):
        h, w = img.shape[: 2]
        pad_h = max(self.crop_size - h, 0)
        pad_w = max(self.crop_size - w, 0)
        img = np.pad(img, ((0, pad_h), (0, pad_w), (0, 0)), 'constant')
        mask = np.pad(mask, ((0, pad_h), (0, pad_w)), 'constant', constant_values=self.ignore_label)
        return img, mask

    def __call__(self, img, mask):
        assert img.size == mask.size

        w, h = img.size
        long_size = max(h, w)

        img = np.array(img)
        mask = np.array(mask)

        if long_size > self.crop_size:
            stride = int(math.ceil(self.crop_size * self.stride_rate))
            h_step_num = int(math.ceil((h - self.crop_size) / float(stride))) + 1
            w_step_num = int(math.ceil((w - self.crop_size) / float(stride))) + 1
            img_sublist, mask_sublist = [], []
            for yy in range(h_step_num):
                for xx in range(w_step_num):
                    sy, sx = yy * stride, xx * stride
                    ey, ex = sy + self.crop_size, sx + self.crop_size
                    img_sub = img[sy: ey, sx: ex, :]
                    mask_sub = mask[sy: ey, sx: ex]
                    img_sub, mask_sub = self._pad(img_sub, mask_sub)
                    img_sublist.append(Image.fromarray(img_sub.astype(np.uint8)).convert('RGB'))
                    mask_sublist.append(Image.fromarray(mask_sub.astype(np.uint8)).convert('P'))
            return img_sublist, mask_sublist
        else:
            img, mask = self._pad(img, mask)
            img = Image.fromarray(img.astype(np.uint8)).convert('RGB')
            mask = Image.fromarray(mask.astype(np.uint8)).convert('P')
            return img, mask


class SlidingCrop(object):
    def __init__(self, crop_size, stride_rate, ignore_label):
        self.crop_size = crop_size
        self.stride_rate = stride_rate
        self.ignore_label = ignore_label

    def _pad(self, img, mask):
        h, w = img.shape[: 2]
        pad_h = max(self.crop_size - h, 0)
        pad_w = max(self.crop_size - w, 0)
        img = np.pad(img, ((0, pad_h), (0, pad_w), (0, 0)), 'constant')
        mask = np.pad(mask, ((0, pad_h), (0, pad_w)), 'constant', constant_values=self.ignore_label)
        return img, mask, h, w

    def __call__(self, img, mask):
        assert img.size == mask.size

        w, h = img.size
        long_size = max(h, w)

        img = np.array(img)
        mask = np.array(mask)

        if long_size > self.crop_size:
            stride = int(math.ceil(self.crop_size * self.stride_rate))
            h_step_num = int(math.ceil((h - self.crop_size) / float(stride))) + 1
            w_step_num = int(math.ceil((w - self.crop_size) / float(stride))) + 1
            img_slices, mask_slices, slices_info = [], [], []
            for yy in range(h_step_num):
                for xx in range(w_step_num):
                    sy, sx = yy * stride, xx * stride
                    ey, ex = sy + self.crop_size, sx + self.crop_size
                    img_sub = img[sy: ey, sx: ex, :]
                    mask_sub = mask[sy: ey, sx: ex]
                    img_sub, mask_sub, sub_h, sub_w = self._pad(img_sub, mask_sub)
                    img_slices.append(Image.fromarray(img_sub.astype(np.uint8)).convert('RGB'))
                    mask_slices.append(Image.fromarray(mask_sub.astype(np.uint8)).convert('P'))
                    slices_info.append([sy, ey, sx, ex, sub_h, sub_w])
            return img_slices, mask_slices, slices_info
        else:
            img, mask, sub_h, sub_w = self._pad(img, mask)
            img = Image.fromarray(img.astype(np.uint8)).convert('RGB')
            mask = Image.fromarray(mask.astype(np.uint8)).convert('P')
            return [img], [mask], [[0, sub_h, 0, sub_w, sub_h, sub_w]]


# --------------------------------------- my transforms ---------------------------------------

class ToTensor(object):
    def __call__(self, images, masks):
        return tuple(F.to_tensor(image) for image in images), \
               tuple(torch.from_numpy(np.array(mask)).float() for mask in masks)


class ScaleFixedWidth(object):
    def __init__(self, width):
        self.width = width

    def __call__(self, image, mask):
        assert image.size == mask.size
        w, h = image.size

        tw = self.width
        th = round(tw * h / w)
        tsize = (tw, th)

        return image.resize(tsize, Image.BILINEAR), mask.resize(tsize, Image.NEAREST)


class ColorJitter(object):
    def __init__(self, brightness=0, contrast=0, saturation=0, hue=0):
        self.t = transforms.ColorJitter(brightness, contrast, saturation, hue)

    def __call__(self, images, masks):
        return (self.t(images[0]),) + images[1:], masks


def get_x_center(mask):
    xs = np.argwhere(np.array(mask).astype(np.bool))[:, 1]
    if len(xs):
        return int(xs.mean())
    return mask.size[0] // 2


class ParametersRandomCrop(object):
    def __init__(self, top_random, min_width, max_width, w2h):  # w * w2h = h
        self.top_random = top_random
        self.min_width = min_width
        self.max_width = max_width
        self.w2h = w2h

    def __call__(self, image, mask):
        assert image.size == mask.size

        a = random.random()
        tw = self.min_width + round(a * (self.max_width - self.min_width))
        th = round(tw * self.w2h)

        x_center = get_x_center(mask)
        x = random.randint(int(x_center - tw * 2 / 3), int(x_center - tw * 1 / 3))
        y = random.randint(*self.top_random)

        return image.crop((x, y, x + tw, y + th)), mask.crop((x, y, x + tw, y + th))


class ParametersCrop(object):
    def __init__(self, width, w2h):  # w * w2h = h
        self.width = width
        self.w2h = w2h

    def __call__(self, image, mask):
        assert image.size == mask.size

        tw = self.width
        th = round(tw * self.w2h)

        x_center = get_x_center(mask)
        x = int(x_center - tw / 2)
        y = 0

        return image.crop((x, y, x + tw, y + th)), mask.crop((x, y, x + tw, y + th))


# class RandomBackground(object):
#     def __init__(self, scene=0., pure=0.):
#         self.scene = scene
#         self.pure = pure

#     def __call__(self, images, masks):
#         if len(images) > 1:
#             image, alpha = images[:2]

#             random_ = random.uniform(0, 1)
#             background = None
#             if random_ < self.pure:
#                 background = random_pure()
#             elif random_ < self.pure + self.scene:
#                 scene_root = r'/root/data/background'
#                 background = cv2.imread(os.path.join(scene_root, random.choice(os.listdir(scene_root))))
#                 background = utils.resize(background, target_w=int(image.size[0] * 1.5))
#                 background = utils.random_crop(background, image.size[::-1])

#             if background is not None:
#                 image = utils.merge(np.array(image)[..., ::-1], np.array(alpha), background)
#                 images = (Image.fromarray(image[..., ::-1]),) + images[1:]

#         return images, masks


def random_pure():
    return (
        random.randint(0, 255),
        random.randint(0, 255),
        random.randint(0, 255),
    )
